#pragma once

#include <iostream>
#include <stdexcept>
#include <string>
#include <memory>
#include <thread>
#include <exception>
#include <cstring>
#include <cerrno>
#include <sys/types.h>
#if defined (__linux__) || defined (__APPLE__)
#    include <sys/socket.h>
#    include <arpa/inet.h>
#    include <netinet/in.h>
#    include <netdb.h>
#    include <unistd.h>
#elif defined (_WIN32) || defined (_WIN64)
#    pragma warning(disable: 4996) // disable vs unsafed warnings
#    include <winsock.h> //marco ERROR defined in wingdi.h may conflicting with your code, define NOGDI can solve it
#    pragma comment(lib,"ws2_32.lib") // cmdline need -lwsock32 -lWs2_32
#endif

#if defined (_WIN32) || defined (_WIN64)
  typedef int socklen_t;
  typedef size_t ssize_t;
#endif

namespace inet
{
struct api
{
    enum protocol
    {
        udp = SOCK_DGRAM,
        tcp = SOCK_STREAM,
    };

    enum iporhost
    {
        useip,
        usehost,
    };

    static int Socket(protocol proto, bool reuse_addr = true)
    {
#if defined (WIN32) || defined (WIN64)
		WSADATA wsaData;
        if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0)
            throw std::runtime_error("winsock init error");
#endif

        int fd = socket(AF_INET, proto, 0);
        if (fd < 0)
            throw std::runtime_error(std::string("socket failed: ") + strerror(errno));

#if defined (__linux__) || defined (__APPLE__)
        int opt = 1;
#elif defined (WIN32) || defined (WIN64)
        char opt = 1;
#endif
        if (reuse_addr && setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt)) < 0)
            throw std::runtime_error(std::string("setsockopt failed: ") + strerror(errno));

        return fd;
    }

    static void Bind(int sock, const std::string& ip, uint16_t port)
    {
        struct sockaddr_in local;
        memset(&local, 0, sizeof(local));

        local.sin_family      = AF_INET;
        local.sin_addr.s_addr = inet_addr(ip.c_str());
        local.sin_port        = htons(port);

        if (bind(sock, (struct sockaddr*)&local, sizeof(local)) < 0)
            throw std::runtime_error(std::string("bind error: ") + strerror(errno));
    }

    static void Bind(int sock, uint16_t port)
    {
        struct sockaddr_in local;
        memset(&local, 0, sizeof(local));

        local.sin_family      = AF_INET;
        local.sin_addr.s_addr = INADDR_ANY;
        local.sin_port        = htons(port);

        if (bind(sock, (struct sockaddr*)&local, sizeof(local)) < 0)
            throw std::runtime_error(std::string("bind error: ") + strerror(errno));
    }

    static void Listen(int sock, int backlog = 12)
    {
        if (listen(sock, backlog) < 0)
            throw std::runtime_error(std::string("listen error: ") + strerror(errno));
    }

    static void Connect(int sock, iporhost ioh, const std::string& iporhost, uint16_t port, int trytime = 1)
    {
        struct hostent* host = nullptr;
        if (ioh == usehost)
        {
            host = gethostbyname(iporhost.c_str());
            if (!host) throw std::runtime_error("get host by name failed");
        }

        struct sockaddr_in peer;
        memset(&peer, 0, sizeof(peer));

        peer.sin_family = AF_INET;
        if (ioh == useip) peer.sin_addr.s_addr = inet_addr(iporhost.c_str());
        else if (ioh == usehost) peer.sin_addr = *(struct in_addr*)host->h_addr;
        peer.sin_port = htons(port);

        while ((trytime-- > 0) && (connect(sock, (struct sockaddr*)&peer, sizeof(peer)) < 0))
        {
            if (trytime == 0) throw std::runtime_error(std::string("connect error: ") + strerror(errno));
            std::this_thread::sleep_for(std::chrono::seconds(1));
        }
    }

    static int Accept(int sock, std::string* ip = nullptr, uint16_t* port = nullptr)
    {
        struct sockaddr_in peer;
        socklen_t len = sizeof(peer);
        memset(&peer, 0, len);

        int fd = accept(sock, (struct sockaddr*)&peer, &len);

        if (ip)   *ip   = inet_ntoa(peer.sin_addr);
        if (port) *port = ntohs(peer.sin_port);

        if (fd < 0)
        {
            if (errno == EAGAIN || errno == EWOULDBLOCK) return fd;
            else throw std::runtime_error(std::string("accept error: ") + strerror(errno));
        }

        return fd;
    }

    static ssize_t Recv(int sock, std::string* msg, size_t maxlen)
    {
        msg->clear();
        std::unique_ptr<char[]> buf(new char[maxlen + 1]{0});

        ssize_t s = recv(sock, buf.get(), maxlen, 0);
        if (s > 0)
        {
            buf[s] = 0;
            *msg = buf.get();
        }
        return s;
    }

    static ssize_t Recv(int sock, char* buf, size_t maxlen)
    {
        ssize_t s = recv(sock, buf, maxlen, 0);
        if (s > 0)
        {
            if (s < maxlen)
            {
                buf[s] = 0;
            }
        }
        return s;
    }

    static ssize_t Recvaline(int sock, std::string* msg, bool nolf = true)
    {
        msg->clear();
        char ch;

        while (true)
        {
            ssize_t s = recv(sock, &ch, 1, 0);

            if (s <= 0)
                break;

            if (nolf)
            {
                if (ch == '\r')
                    recv(sock, &ch, 1, 0);
                if (ch == '\n')
                    break;
            }

            msg->push_back(ch);
            if (ch == '\n')
                break;
        }

        return msg->size();
    }

#if defined (__linux__) || defined (__APPLE__)
    static ssize_t Recvall(int sock, std::string* msg)
    {
        msg->clear();
        std::unique_ptr<char[]> buf(new char[1024]{0});

        while (true)
        {

            ssize_t s = recv(sock, buf.get(), 1024, MSG_DONTWAIT);

            if (s > 0)
            {
                buf[s] = 0;
                *msg += buf.get();
            }
            else if (s <= 0)
            {
                break;
            }
        }

        return msg->size();
    }
#endif

    static ssize_t Send(int sock, const std::string& msg)
    {
        return send(sock, msg.c_str(), msg.size(), 0);
    }

    static ssize_t Recvfrom(int sock, std::string* msg, size_t maxlen, std::string* ip = nullptr, uint16_t* port = nullptr)
    {
        msg->clear();
        std::unique_ptr<char[]> buf(new char[maxlen]{0});

        struct sockaddr_in peer;
        socklen_t len = sizeof(peer);
        memset(&peer, 0, len);
        ssize_t s = recvfrom(sock, buf.get(), maxlen, 0, (struct sockaddr*)&peer, &len);

        if (s > 0)
        {
            buf[s] = 0;
            *msg = buf.get();
        }
        if (ip)   *ip   = inet_ntoa(peer.sin_addr);
        if (port) *port = ntohs(peer.sin_port);

        return s;
    }

    static ssize_t Recvfrom(int sock, char* buf, size_t maxlen, std::string* ip = nullptr, uint16_t* port = nullptr)
    {
        struct sockaddr_in peer;
        socklen_t len = sizeof(peer);
        memset(&peer, 0, len);
        ssize_t s = recvfrom(sock, buf, maxlen, 0, (struct sockaddr*)&peer, &len);

        if (s > 0)
        {
            buf[s] = 0;
        }
        if (ip)   *ip   = inet_ntoa(peer.sin_addr);
        if (port) *port = ntohs(peer.sin_port);

        return s;
    }

    static ssize_t Sendto(int sock, const std::string& msg, const std::string& ip, uint16_t port)
    {
        struct sockaddr_in peer;
        memset(&peer, 0, sizeof(peer));

        peer.sin_family      = AF_INET;
        peer.sin_addr.s_addr = inet_addr(ip.c_str());
        peer.sin_port        = htons(port);

        ssize_t s = sendto(sock, msg.c_str(), msg.size(), 0, (struct sockaddr*)&peer, sizeof(peer));
        return s;
    }

    static void Close(int fd)
    {
#if defined (__linux__) || defined (__APPLE__)
        close(fd);
#elif defined (WIN32) || defined (WIN64)
        closesocket(fd);
        WSACleanup();
#endif
    }
};

namespace tcp
{
class server
{
public:
    server(uint16_t port, int backlog = 12)
        : _sock(-1), _port(port), _backlog(backlog)
    {
        __init__();
    }

    server(const std::string& ip, uint16_t port, int backlog = 12)
        : _sock(-1), _ip(ip), _port(port), _backlog(backlog)
    {
        __init__();
    }

    int accept(std::string* cip = nullptr, uint16_t* cport = nullptr) { return inet::api::Accept(_sock, cip, cport); }

    ssize_t recv     (int sock, std::string* msg, size_t maxlen)    { return inet::api::Recv(sock, msg, maxlen);    }
    ssize_t recv     (int sock,        char* buf, size_t maxlen)    { return inet::api::Recv(sock, buf, maxlen);    }
    ssize_t recvaline(int sock, std::string* msg, bool nolf = true) { return inet::api::Recvaline(sock, msg, nolf); }
#if defined (__linux__) || defined (__APPLE__)
    ssize_t recvall  (int sock, std::string* msg)                   { return inet::api::Recvall(sock, msg);         }
#endif

    ssize_t send(int sock, const std::string& msg) { return inet::api::Send(sock, msg); }

    ~server() { inet::api::Close(_sock); }

public:
    int         sock()    const { return _sock;    }
    std::string ip()      const { return _ip;      }
    uint16_t    port()    const { return _port;    }
    int         backlog() const { return _backlog; }

private:
    void __init__()
    {
        _sock = inet::api::Socket(inet::api::tcp);

        if (_ip.empty()) inet::api::Bind(_sock, _port);
        else             inet::api::Bind(_sock, _ip, _port);

        inet::api::Listen(_sock, _backlog);
    }

protected:
    int _sock;
    std::string _ip;
    uint16_t _port;
    int _backlog;
};

class client
{
public:
    client(inet::api::iporhost ioh, const std::string& svr_iporhost, uint16_t svr_port, int trytime = 1)
        : _sock(-1), _siporhost(svr_iporhost), _sport(svr_port), _trytime(trytime)
    {
        _sock = inet::api::Socket(inet::api::tcp);
        inet::api::Connect(_sock, ioh, _siporhost, _sport, _trytime);
    }

    ssize_t send(int sock, const std::string& msg) { return inet::api::Send( sock, msg); }
    ssize_t send(          const std::string& msg) { return inet::api::Send(_sock, msg); }

    ssize_t recv     (int sock, std::string* msg, size_t maxlen)    { return inet::api::Recv( sock, msg, maxlen);    }
    ssize_t recv     (          std::string* msg, size_t maxlen)    { return inet::api::Recv(_sock, msg, maxlen);    }
    ssize_t recv     (int sock,        char* buf, size_t maxlen)    { return inet::api::Recv( sock, buf, maxlen);    }
    ssize_t recv     (                 char* buf, size_t maxlen)    { return inet::api::Recv(_sock, buf, maxlen);    }
    ssize_t recvaline(int sock, std::string* msg, bool nolf = true) { return inet::api::Recvaline( sock, msg, nolf); }
    ssize_t recvaline(          std::string* msg, bool nolf = true) { return inet::api::Recvaline(_sock, msg, nolf); }
#if defined (__linux__) || defined (__APPLE__)
    ssize_t recvall  (int sock, std::string* msg)                   { return inet::api::Recvall( sock, msg);         }
    ssize_t recvall  (          std::string* msg)                   { return inet::api::Recvall(_sock, msg);         }
#endif

    ~client() { inet::api::Close(_sock); }

public:
    int         sock()  const { return _sock;      }
    std::string sip()   const { return _siporhost; }
    uint16_t    sport() const { return _sport;     }

protected:
    int _sock;
    std::string _siporhost;
    uint16_t _sport;
    int _trytime;
};
}

namespace udp
{
class server
{
public:
    server(uint16_t port) : _sock(-1), _port(port)
    {
        __init__();
    }

    server(const std::string& ip, uint16_t port) : _sock(-1), _ip(ip), _port(port)
    {
        __init__();
    }

    ssize_t sendto(const std::string& msg, const std::string& cip, uint16_t cport)
    {
        return inet::api::Sendto(_sock, msg, cip, cport);
    }

    ssize_t recvfrom(std::string* msg, size_t maxlen, std::string* cip = nullptr, uint16_t* cport = nullptr)
    {
        return inet::api::Recvfrom(_sock, msg, maxlen, cip, cport);
    }
    ssize_t recvfrom(       char* buf, size_t maxlen, std::string* cip = nullptr, uint16_t* cport = nullptr)
    {
        return inet::api::Recvfrom(_sock, buf, maxlen, cip, cport);
    }

    ~server() { inet::api::Close(_sock); }

public:
    int         sock()    const { return _sock;    }
    std::string ip()      const { return _ip;      }
    uint16_t    port()    const { return _port;    }

private:
    void __init__()
    {
        _sock = inet::api::Socket(inet::api::udp);

        if (_ip.empty()) inet::api::Bind(_sock,      _port);
        else             inet::api::Bind(_sock, _ip, _port);
    }

protected:
    int _sock;
    std::string _ip;
    uint16_t _port;
};

class client
{
public:
    client(const std::string& svr_ip, uint16_t svr_port)
        : _sock(-1), _sip(svr_ip), _sport(svr_port)
    {
        _sock = inet::api::Socket(inet::api::udp);
    }

    ssize_t sendto(const std::string& msg) { return inet::api::Sendto(_sock, msg, _sip, _sport); }

    ssize_t recvfrom(std::string* msg, size_t maxlen) { return inet::api::Recvfrom(_sock, msg, maxlen); }
    ssize_t recvfrom(       char* buf, size_t maxlen) { return inet::api::Recvfrom(_sock, buf, maxlen); }

    ~client() { inet::api::Close(_sock); }

public:
    int         sock()  const { return _sock;  }
    std::string sip()   const { return _sip;   }
    uint16_t    sport() const { return _sport; }

protected:
    int _sock;
    std::string _sip;
    uint16_t _sport;
};
}

}
